import cv2
import numpy as np
import os
import argparse
import logging
from pathlib import Path

from edge_detection import apply_canny_edge_detection
from kmeans_clustering import apply_kmeans_clustering

# 设置日志
def setup_logger():
    logger = logging.getLogger('ImageProcessor')
    if not logger.handlers:
        logger.setLevel(logging.INFO)
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        logger.addHandler(handler)
    return logger

logger = setup_logger()

def process_image(image_path, output_dir=None, n_clusters=None):
    """
    处理单张图像，包括边缘检测和颜色聚类
    
    Args:
        image_path: 输入图像路径
        output_dir: 输出目录，如果为None则使用图像所在目录
        n_clusters: 指定聚类数量，如果为None则自动确定
    
    Returns:
        成功返回True，失败返回False
    """
    try:
        # 读取图像
        logger.info(f"读取图像: {image_path}")
        image = cv2.imread(str(image_path))
        if image is None:
            logger.error(f"无法读取图像: {image_path}")
            return False
        
        # 创建输出目录
        if output_dir is None:
            output_dir = Path(image_path).parent / f"{Path(image_path).stem}_output"
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"输出目录: {output_dir}")
        
        # 保存原始图像
        original_output_path = os.path.join(output_dir, "1_原图.png")
        cv2.imwrite(original_output_path, image)
        
        # 执行边缘检测
        logger.info("执行Canny边缘检测...")
        edges = apply_canny_edge_detection(image)
        edge_output_path = os.path.join(output_dir, "2_边缘检测.png")
        cv2.imwrite(edge_output_path, edges)
        
        # 执行颜色聚类
        logger.info("执行K-means++聚类...")
        result = apply_kmeans_clustering(image, n_clusters)
        if result is None:
            logger.error("聚类处理失败")
            return False
            
        clustered_image, layer_results, layer_masks = result
        
        # 保存聚类结果
        clustered_output_path = os.path.join(output_dir, "3_聚类结果.png")
        cv2.imwrite(clustered_output_path, clustered_image)
        
        # 保存每个图层和掩码
        logger.info(f"保存 {len(layer_results)} 个聚类图层...")
        for i, (layer, mask) in enumerate(zip(layer_results, layer_masks)):
            layer_output_path = os.path.join(output_dir, f"4_{i+1}_图层.png")
            mask_output_path = os.path.join(output_dir, f"5_{i+1}_掩码.png")
            cv2.imwrite(layer_output_path, layer)
            cv2.imwrite(mask_output_path, mask)
        
        logger.info("图像处理完成")
        return True
        
    except Exception as e:
        logger.error(f"处理图像时发生错误: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return False

def batch_process(input_folder, output_folder=None, n_clusters=None):
    """
    批量处理文件夹中的所有图像
    
    Args:
        input_folder: 输入文件夹路径
        output_folder: 输出文件夹路径，如果为None则在输入文件夹中创建output子文件夹
        n_clusters: 指定聚类数量，如果为None则自动确定
    
    Returns:
        成功处理的图像数量
    """
    input_folder = Path(input_folder)
    if not input_folder.exists() or not input_folder.is_dir():
        logger.error(f"输入文件夹不存在或不是有效目录: {input_folder}")
        return 0
        
    if output_folder is None:
        output_folder = input_folder / "output"
    os.makedirs(output_folder, exist_ok=True)
    
    # 获取所有图像文件
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']
    image_files = []
    for ext in image_extensions:
        image_files.extend(list(input_folder.glob(f"*{ext}")))
        image_files.extend(list(input_folder.glob(f"*{ext.upper()}")))
    
    if not image_files:
        logger.warning(f"输入文件夹 {input_folder} 中没有找到图像文件")
        return 0
    
    logger.info(f"找到 {len(image_files)} 个图像文件，开始批量处理...")
    
    success_count = 0
    for i, image_path in enumerate(image_files):
        logger.info(f"处理第 {i+1}/{len(image_files)} 个图像: {image_path.name}")
        image_output_dir = output_folder / image_path.stem
        if process_image(image_path, image_output_dir, n_clusters):
            success_count += 1
    
    logger.info(f"批量处理完成，成功处理 {success_count}/{len(image_files)} 个图像")
    return success_count

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="图像边缘检测和颜色聚类工具")
    parser.add_argument("--input", "-i", type=str, required=True, help="输入图像或文件夹路径")
    parser.add_argument("--output", "-o", type=str, default=None, help="输出目录路径")
    parser.add_argument("--clusters", "-c", type=int, default=None, help="指定聚类数量（默认自动确定）")
    parser.add_argument("--batch", "-b", action="store_true", help="批量处理模式")
    
    args = parser.parse_args()
    
    if args.batch:
        batch_process(args.input, args.output, args.clusters)
    else:
        process_image(args.input, args.output, args.clusters) 